import os
import sys
import platform
import torch
from pathlib import Path
from typing import Dict, Any, Optional
from .exceptions import InvalidDeviceError


class Environment:
    """增强版环境配置类（整合原环境+配置功能）"""
    # __init__方法用于初始化Environment对象。它接受一个可选的参数config_overrides，这是一个字典，用于覆盖默认的配置。
    def __init__(self, config_overrides: Optional[Dict[str, Any]] = None):
        # 基础路径配置,这里设置了项目的基础路径project_root。
        self.project_root = Path(__file__).parent.parent
        # 并调用_setup_path方法来配置Python路径
        self._setup_path()

        # 设备配置
        self.is_windows = platform.system() == 'Windows'
        self._configure_device()

        # 训练参数（带默认值）
        self.batch_size = 64 if self.is_windows else 128 # 如果是Windows系统，batch_size为64，否则为128
        self.epochs = 10 # 默认训练轮数为10
        self.lr = 0.001 # 默认学习率为0.001,学习率（Learning Rate，简称 lr）是深度学习中优化器更新模型参数时的步长系数，控制参数调整的幅度。
        # 权重衰减（Weight Decay）是一种正则化技术，通过在损失函数中添加模型参数的 L2 范数惩罚项，限制参数值过大，从而防止模型过拟合。
        self.weight_decay = 1e-5 # 默认权重衰减为1e-5

        # 覆盖用户自定义配置
        if config_overrides:
            for key, value in config_overrides.items():
                setattr(self, key, value)

        # 创建目录
        self.checkpoint_dir = self.project_root / "checkpoints"
        self.log_dir = self.project_root / "logs"
        self._create_dirs()

    def _setup_path(self):
        """配置Python路径"""
        if str(self.project_root) not in sys.path:
            sys.path.insert(0, str(self.project_root))

    def _configure_device(self):
        """设备配置逻辑（兼容原设计）"""
        if os.getenv('FORCE_CPU') == '1':
            self._forced_cpu = True
            os.environ['CUDA_VISIBLE_DEVICES'] = '-1'
        else:
            self._forced_cpu = False
            if self.is_windows:
                os.environ['CUDA_VISIBLE_DEVICES'] = '-1'
            else:
                os.environ['CUDA_VISIBLE_DEVICES'] = os.getenv('CUDA_VISIBLE_DEVICES', '0')

    def _create_dirs(self):
        """创建必要目录"""
        self.checkpoint_dir.mkdir(parents=True, exist_ok=True)
        self.log_dir.mkdir(parents=True, exist_ok=True)

    @property
    def device(self) -> torch.device:
        """获取当前设备（增强版）"""
        if self._forced_cpu:
            return torch.device('cpu')

        if self.is_windows:
            return torch.device('cpu')

        if not torch.cuda.is_available():
            return torch.device('cpu')

        device_id = int(os.environ['CUDA_VISIBLE_DEVICES'].split(',')[0])
        if device_id >= torch.cuda.device_count():
            raise InvalidDeviceError(
                requested=f"cuda:{device_id}",
                available=f"{torch.cuda.device_count()} GPU(s)"
            )
        return torch.device(f'cuda:{device_id}')

    def get_device_info(self) -> Dict[str, Any]:
        """获取设备详细信息"""
        device = self.device
        info = {
            'device_type': device.type,
            'device_index': device.index if device.index is not None else 0
        }

        if device.type == 'cuda':
            info.update({
                'gpu_name': torch.cuda.get_device_name(device),
                'cuda_version': torch.version.cuda,
                'total_memory_gb': torch.cuda.get_device_properties(device).total_memory / 1024 ** 3
            })
        return info

    # 这个方法返回当前设备的详细信息，包括设备类型、设备索引、GPU名称、CUDA版本和总内存。
    def to_dict(self) -> Dict[str, Any]:
        """导出配置字典"""
        return {
            'device': str(self.device),
            'batch_size': self.batch_size,
            'epochs': self.epochs,
            'project_root': str(self.project_root)
        }